from ultralytics import YOLO

#导入中断的模型继续训练
model = YOLO("./runs/detect/train/weights/last.pt")  # 加载预训练模型（建议用于训练）

if __name__ == '__main__':
    # 单GPU训练
    model.train(
        resume=True,
        data="dataset.yaml",  # 你的数据集配置文件
        # 数据增强参数
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=10.0,  # 允许±10度旋转
        fliplr=0.5,  # 50%概率左右翻转
        mosaic=0.8,  # 80%概率使用Mosaic增强
        mixup=0.2  # 20%概率使用MixUp增强
    ) # 训练模型
    # 多GPU训练
    # results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])
    metrics = model.val()  # 在验证集上评估模型性能
    success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
    results = model("./dataset/test/images/00002.png")  # 对图像进行预测
    # Process results list
    for result in results:
        boxes = result.boxes  # Boxes object for bounding box outputs
        masks = result.masks  # Masks object for segmentation masks outputs
        keypoints = result.keypoints  # Keypoints object for pose outputs
        probs = result.probs  # Probs object for classification outputs
        result.show()  # display to screen